import os
import numpy as np
import rasterio
from rasterio.plot import show
import matplotlib.pyplot as plt
from scipy.stats import f
from tqdm import trange

plt.rcParams['font.sans-serif'] = ['Noto Sans CJK JP']  # 选择简体中文
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 文件路径
input_path = './correlation_outputs/'  # 输入文件夹路径
output_path = './correlation_outputs/'  # 输出文件夹路径

# 确保输出文件夹存在
os.makedirs(output_path, exist_ok=True)

# 加载偏相关系数 tif 文件
def load_tif(file_path):
    """加载 tif 文件并返回数据和元数据"""
    with rasterio.open(file_path) as src:
        data = src.read(1)
        meta = src.meta
    return data, meta

# 计算复相关系数及其显著性 F 检验
def calculate_multiple_correlation_with_ftest(r_xy, r_xz, r_yz, n):
    """计算复相关系数及其显著性 F 检验"""
    # 检查分母是否为零
    denominator = 1 - r_yz**2
    if denominator == 0:
        return np.nan, np.nan

    # 计算复相关系数
    numerator = (1 - r_xy**2) * (1 - r_xz**2)
    r_multiple = np.sqrt(1 - numerator / denominator)

    # 确保复相关系数在 [0, 1] 范围内
    if r_multiple < 0 or r_multiple > 1:
        return np.nan, np.nan

    # 计算 F 统计量
    k = 2  # 自变量数量
    f_stat = (r_multiple**2 / (1 - r_multiple**2)) * ((n - k - 1) / k)

    # 计算 p 值
    #p_value = 1 - f.cdf(f_stat, k, n - k - 1)

    return r_multiple, f_stat#p_value

# 加载降雨量和气温的偏相关系数 tif 文件
rainfall_corr, meta = load_tif(os.path.join(input_path, 'partial_correlation_rainfall.tif'))
temperature_corr, _ = load_tif(os.path.join(input_path, 'partial_correlation_temperature.tif'))

# 假设样本数量（时间序列长度）
n = 21  # 例如 21 年的数据

# 初始化结果数组
rows, cols = rainfall_corr.shape
multiple_corr = np.full((rows, cols), np.nan)
pvalue_multiple = np.full((rows, cols), np.nan)

# 计算每个像素的复相关系数和显著性
for i in trange(rows, desc="计算复相关系数"):
    for j in range(cols):
        r_xy = rainfall_corr[i, j]
        r_xz = temperature_corr[i, j]
        r_yz = 0  # 假设降雨量和气温之间的相关系数为 0（如果有数据可以替换）
        
        # 跳过无效值
        if np.isnan(r_xy) or np.isnan(r_xz):
            continue
        
        # 计算复相关系数和显著性
        multiple_corr[i, j], pvalue_multiple[i, j] = calculate_multiple_correlation_with_ftest(
            r_xy, r_xz, r_yz, n
        )

# 保存复相关系数和显著性结果
meta.update(dtype='float32', count=1, nodata=np.nan)

with rasterio.open(os.path.join(output_path, 'multiple_correlation.tif'), 'w', **meta) as dst:
    dst.write(multiple_corr.astype('float32'), 1)

with rasterio.open(os.path.join(output_path, 'F_value_multiple.tif'), 'w', **meta) as dst:
    dst.write(pvalue_multiple.astype('float32'), 1)

# 可视化复相关系数和显著性分布
plt.figure(figsize=(15, 8))

# 创建 1x2 的子图
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 8))

# 绘制复相关系数分布图
im1 = ax1.imshow(multiple_corr, 
                 cmap='RdBu_r',
                 vmin=0, 
                 vmax=1)
ax1.set_title('复相关系数空间分布')
plt.colorbar(im1, ax=ax1, label='复相关系数')

# 绘制显著性分布图
significance_multiple = np.where(pvalue_multiple < 0.05, 1, 0)
im2 = ax2.imshow(significance_multiple, 
                 cmap='binary',
                 vmin=0,
                 vmax=1)
ax2.set_title('复相关系数显著性分布\n(黑色区域 p < 0.05)')
plt.colorbar(im2, ax=ax2, label='显著性')

plt.tight_layout()
plt.savefig(os.path.join(output_path, 'multiple_correlation_with_significance.png'), 
            dpi=300, 
            bbox_inches='tight')
plt.show()

# 打印统计信息
print(f"复相关系数显著相关区域占比: {np.sum(significance_multiple) / significance_multiple.size * 100:.2f}%")